/*
  Author: Chen Haodong (shinkunyachen@outlook.com)
  -----
  Copyright (c)2020 - present Chen Haodong. All Rights Reserved.
 */

#include "checkers/BranchChecker/LCFG.h"
#include <utility>

void LCFGStatement::dump(const clang::ASTContext &ctx)
{
    llvm::errs() << "Statement " << _id << " in BB" << _basic_block->id() << '\n';
    llvm::errs() << "in Plane:";
    for (auto p : _planes)
    {
        llvm::errs() << ' ' << p->id();
    }
    llvm::errs() << '\n';
    _stmt->dumpPretty(ctx);
    llvm::errs() << '\n';
}

void LCFGPlane::dump(const clang::ASTContext &ctx)
{
    llvm::errs() << "----------Plane(" << _id << ")----------\n";
    if (_dominator)
    {
        _dominator->dump(ctx);
    }
    for (auto i : _elements)
    {
        i->dump(ctx);
    }
    llvm::errs() << "----------Plane End----------\n";
}

void LCFGBasicBlock::dump(const clang::ASTContext &ctx)
{
    llvm::errs() << "----------BasicBlock(" << _id << ")----------\n";
    for (auto i : _element_sequence)
    {
        i->dump(ctx);
    }
    llvm::errs() << "Preds:";
    for (auto i : _in_blocks)
    {
        llvm::errs() << ' ' << i->id();
    }
    llvm::errs() << "\nSuccs:";
    for (auto i : _out_blocks)
    {
        llvm::errs() << ' ' << i->id();
    }
    llvm::errs() << '\n';
    llvm::errs() << "----------BasicBlock End----------\n";
}

void LCFG::build()
{

    _main_plane = new_plane();
    BuildContext ctx;
    ctx.exit_jmp_target = new LCFGJumpTarget;
    ctx.plane_stack.push_back(_main_plane);
    clang::CompoundStmt *body =
        llvm::dyn_cast_or_null<clang::CompoundStmt>(_function_decl->getBody());
    if (!body)
        return;
    dispatch_element_generation(body, ctx);
    std::vector<std::pair<LCFGJump *, LCFGBasicBlock *>>
        jmp_list;
    std::unordered_map<LCFGJumpTarget *, LCFGBasicBlock *>
        jmp_target_map;
    _entry_block = new_basic_block();
    _exit_block = new_basic_block();
    LCFGBasicBlock *cur_block = new_basic_block();
    link_basic_block(_entry_block, cur_block);
    jmp_target_map.insert(std::make_pair(ctx.exit_jmp_target, _exit_block));
    bool block_end = false;
    for (auto i : _elements)
    {
        switch (i->kind())
        {
        case LCFGElement::Statement:
            if (block_end)
            {
                cur_block = new_basic_block();
                block_end = false;
            }
            add_element_to_basic_block(i, cur_block);
            break;
        case LCFGElement::Jump:
            block_end = true;
            jmp_list.push_back(std::make_pair(static_cast<LCFGJump *>(i), cur_block));
            break;
        case LCFGElement::JumpTarget:
            cur_block = new_basic_block();
            block_end = false;
            jmp_target_map.insert(
                std::make_pair(static_cast<LCFGJumpTarget *>(i), cur_block));
            break;
        }
    }
    if (_function_decl->getReturnType()->isVoidType())
    {
        link_basic_block(cur_block, _exit_block);
    }
    for (auto &i : jmp_list)
    {
        link_basic_block(i.second, jmp_target_map.find(i.first->target)->second);
    }
    remove_illegal_successor();
    remove_useless_element();
    remove_empty_basic_block();
    remove_dead_basic_block();
    remove_useless_element();
    for (size_t i = 0, sz = _elements.size(); i < sz; ++i)
    {
        _elements[i]->id(i);
    }
    for (size_t i = 0, sz = _basic_blocks.size(); i < sz; ++i)
    {
        _basic_blocks[i]->id(i);
    }
    delete ctx.exit_jmp_target;
}
void LCFG::dispatch_element_generation(clang::Stmt *stmt, BuildContext &ctx)
{
    if (!stmt)
        return;
    switch (stmt->getStmtClass())
    {
    case clang::Stmt::CompoundStmtClass:
        generate_element(llvm::cast<clang::CompoundStmt>(stmt), ctx);
        break;
    case clang::Stmt::IfStmtClass:
        generate_element(llvm::cast<clang::IfStmt>(stmt), ctx);
        break;
    case clang::Stmt::ForStmtClass:
        generate_element(llvm::cast<clang::ForStmt>(stmt), ctx);
        break;
    case clang::Stmt::WhileStmtClass:
        generate_element(llvm::cast<clang::WhileStmt>(stmt), ctx);
        break;
    case clang::Stmt::DoStmtClass:
        generate_element(llvm::cast<clang::DoStmt>(stmt), ctx);
        break;
    case clang::Stmt::ReturnStmtClass:
        generate_element(llvm::cast<clang::ReturnStmt>(stmt), ctx);
        break;
    case clang::Stmt::BreakStmtClass:
        generate_element(llvm::cast<clang::BreakStmt>(stmt), ctx);
        break;
    case clang::Stmt::ContinueStmtClass:
        generate_element(llvm::cast<clang::ContinueStmt>(stmt), ctx);
        break;
    case clang::Stmt::LabelStmtClass:
        break;
    default:
        generate_element(stmt, ctx);
        break;
    }
}
void LCFG::generate_element(clang::CompoundStmt *stmt, BuildContext &ctx)
{
    for (auto it = stmt->body_begin(), end = stmt->body_end(); it != end; ++it)
    {
        dispatch_element_generation(*it, ctx);
    }
}
void LCFG::generate_element(clang::IfStmt *stmt, BuildContext &ctx)
{
    LCFGStatement *lcfg_stmt = new_element<LCFGStatement>(stmt);
    bind_element_plane(lcfg_stmt, ctx.plane_stack);
    LCFGJump *true_jmp = new_element<LCFGJump>();
    LCFGJump *false_jmp = new_element<LCFGJump>();
    LCFGJumpTarget *true_target = new_element<LCFGJumpTarget>();
    dispatch_element_generation(stmt->getThen(), ctx);
    LCFGJump *next_jmp1 = new_element<LCFGJump>();
    LCFGJumpTarget *false_target = new_element<LCFGJumpTarget>();
    dispatch_element_generation(stmt->getElse(), ctx);
    LCFGJump *next_jmp2 = new_element<LCFGJump>();
    LCFGJumpTarget *next_target = new_element<LCFGJumpTarget>();
    true_jmp->target = true_target;
    false_jmp->target = false_target;
    next_jmp1->target = next_jmp2->target = next_target;
}
void LCFG::generate_element(clang::ForStmt *stmt, BuildContext &ctx)
{
    LCFGPlane *p = new_plane();
    LCFGStatement *lcfg_stmt = new_element<LCFGStatement>(stmt);
    p->dominator(lcfg_stmt);
    ctx.plane_stack.push_back(p);
    ctx.next_jmps_stack.push_back(std::vector<LCFGJump *>{});
    bind_element_plane(lcfg_stmt, ctx.plane_stack);
    LCFGJump *true_jmp = new_element<LCFGJump>();
    LCFGJump *next_jmp = new_element<LCFGJump>();
    LCFGJumpTarget *true_target = new_element<LCFGJumpTarget>();
    dispatch_element_generation(stmt->getBody(), ctx);
    LCFGJump *next_jmp2 = new_element<LCFGJump>();
    LCFGJumpTarget *next_target = new_element<LCFGJumpTarget>();
    true_jmp->target = true_target;
    next_jmp->target = next_jmp2->target = next_target;
    for (auto i : ctx.next_jmps_stack.back())
        i->target = next_target;
    ctx.next_jmps_stack.pop_back();
    ctx.plane_stack.pop_back();
}
void LCFG::generate_element(clang::WhileStmt *stmt, BuildContext &ctx)
{
    LCFGPlane *p = new_plane();
    LCFGStatement *lcfg_stmt = new_element<LCFGStatement>(stmt);
    p->dominator(lcfg_stmt);
    ctx.plane_stack.push_back(p);
    ctx.next_jmps_stack.push_back(std::vector<LCFGJump *>{});
    bind_element_plane(lcfg_stmt, ctx.plane_stack);
    LCFGJump *true_jmp = new_element<LCFGJump>();
    LCFGJump *next_jmp = new_element<LCFGJump>();
    LCFGJumpTarget *true_target = new_element<LCFGJumpTarget>();
    dispatch_element_generation(stmt->getBody(), ctx);
    LCFGJump *next_jmp2 = new_element<LCFGJump>();
    LCFGJumpTarget *next_target = new_element<LCFGJumpTarget>();
    true_jmp->target = true_target;
    next_jmp->target = next_jmp2->target = next_target;
    for (auto i : ctx.next_jmps_stack.back())
        i->target = next_target;
    ctx.next_jmps_stack.pop_back();
    ctx.plane_stack.pop_back();
}
void LCFG::generate_element(clang::DoStmt *stmt, BuildContext &ctx)
{
    LCFGPlane *p = new_plane();
    LCFGStatement *lcfg_stmt = new_element<LCFGStatement>(stmt);
    p->dominator(lcfg_stmt);
    ctx.plane_stack.push_back(p);
    ctx.next_jmps_stack.push_back(std::vector<LCFGJump *>{});
    bind_element_plane(lcfg_stmt, ctx.plane_stack);
    LCFGJump *true_jmp = new_element<LCFGJump>();
    LCFGJump *next_jmp = new_element<LCFGJump>();
    LCFGJumpTarget *true_target = new_element<LCFGJumpTarget>();
    dispatch_element_generation(stmt->getBody(), ctx);
    LCFGJump *next_jmp2 = new_element<LCFGJump>();
    LCFGJumpTarget *next_target = new_element<LCFGJumpTarget>();
    true_jmp->target = true_target;
    next_jmp->target = next_jmp2->target = next_target;
    for (auto i : ctx.next_jmps_stack.back())
        i->target = next_target;
    ctx.next_jmps_stack.pop_back();
    ctx.plane_stack.pop_back();
}
void LCFG::generate_element(clang::ReturnStmt *stmt, BuildContext &ctx)
{
    LCFGStatement *lcfg_stmt = new_element<LCFGStatement>(stmt);
    bind_element_plane(lcfg_stmt, ctx.plane_stack);
    LCFGJump *jmp = new_element<LCFGJump>();
    jmp->target = ctx.exit_jmp_target;
}
void LCFG::generate_element(clang::ContinueStmt *stmt, BuildContext &ctx)
{
    LCFGStatement *lcfg_stmt = new_element<LCFGStatement>(stmt);
    bind_element_plane(lcfg_stmt, ctx.plane_stack);
    LCFGJump *next_jmp = new_element<LCFGJump>();
    ctx.next_jmps_stack.back().push_back(next_jmp);
}
void LCFG::generate_element(clang::BreakStmt *stmt, BuildContext &ctx)
{
    LCFGStatement *lcfg_stmt = new_element<LCFGStatement>(stmt);
    bind_element_plane(lcfg_stmt, ctx.plane_stack);
    LCFGJump *next_jmp = new_element<LCFGJump>();
    ctx.next_jmps_stack.back().push_back(next_jmp);
}
void LCFG::generate_element(clang::Stmt *stmt, BuildContext &ctx)
{
    LCFGStatement *lcfg_stmt = new_element<LCFGStatement>(stmt);
    bind_element_plane(lcfg_stmt, ctx.plane_stack);
}

void LCFG::remove_useless_element()
{
    std::vector<LCFGElement *> elements;
    elements.reserve(_elements.size());
    for (auto i : _elements)
    {
        if (i->kind() == LCFGElement::Statement && i->basic_block())
        {
            elements.push_back(i);
        }
        else
        {
            delete i;
        }
    }
    elements.shrink_to_fit();
    _elements = std::move(elements);
}
void LCFG::remove_empty_basic_block()
{
    std::vector<LCFGBasicBlock *> basic_blocks(_basic_blocks.begin(),
                                               _basic_blocks.begin() + 2);
    basic_blocks.reserve(_basic_blocks.size());
    for (size_t i = 2, sz = _basic_blocks.size(); i < sz; ++i)
    {
        LCFGBasicBlock *bb = _basic_blocks[i];
        if (bb->size() != 0)
        {
            basic_blocks.push_back(bb);
        }
        else
        {
            for (auto j : bb->_in_blocks)
            {
                for (auto k : bb->_out_blocks)
                {
                    link_basic_block(j, k);
                }
            }
            for (auto j : bb->_in_blocks)
            {
                j->erase_succ(bb);
            }
            for (auto j : bb->_out_blocks)
            {
                j->erase_pred(bb);
            }
            delete bb;
        }
    }
    basic_blocks.shrink_to_fit();
    _basic_blocks = std::move(basic_blocks);
}
void LCFG::remove_dead_basic_block()
{
    std::vector<LCFGBasicBlock *> basic_blocks(_basic_blocks.begin(),
                                               _basic_blocks.begin() + 2);
    basic_blocks.reserve(_basic_blocks.size());
    bool remove = false;
    for (size_t i = 2, sz = _basic_blocks.size(); i < sz; ++i)
    {
        remove = false;
        LCFGBasicBlock *bb = _basic_blocks[i];
        if (bb->_in_blocks.size() == 0)
        {
            remove = true;
            for (auto j : bb->_out_blocks)
            {
                j->erase_pred(bb);
            }
            for (auto j : bb->_element_sequence)
            {
                j->basic_block(nullptr);
            }
        }
        if (bb->_out_blocks.size() == 0)
        {
            remove = true;
            for (auto j : bb->_in_blocks)
            {
                j->erase_succ(bb);
            }
            for (auto j : bb->_element_sequence)
            {
                j->basic_block(nullptr);
            }
        }
        if (remove)
        {
            delete bb;
        }
        else
        {
            basic_blocks.push_back(bb);
        }
    }
    basic_blocks.shrink_to_fit();
    _basic_blocks = std::move(basic_blocks);
}
void LCFG::remove_illegal_successor()
{
    std::unordered_set<LCFGBasicBlock *> set({_exit_block});
    for (auto bb : _basic_blocks)
    {
        if (bb->_out_blocks.find(_exit_block) != bb->_out_blocks.end())
        {
            for (auto succ : bb->_out_blocks)
            {
                if (succ != _exit_block)
                {
                    succ->erase_pred(bb);
                }
            }
            bb->_out_blocks = set;
        }
    }
}